import torch
from darknet import Darknet
from dataset import CocoDataset
from torchvision.transforms import transforms
from torch.optim import Adam

transform = transforms.Compose([
    transforms.Resize(size=(416,416)),
    transforms.ToTensor(),
])


device_ids = [0]

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = Darknet(in_chanel=3,use_fc=False).to(device)
#model.load_state_dict(torch.load("./ep69_yolo_sub.ckpt"))
model = torch.nn.DataParallel(model, device_ids=device_ids)
model = model.cuda(device=device_ids[0])

dataset =  CocoDataset("./sub_set/train2014","./sub_set/labels/train2014",to_square=True, transform=transform)
dataloader = torch.utils.data.DataLoader(
    dataset,
    batch_size=8,
    shuffle=True,
    pin_memory=True,
    collate_fn=dataset.collate_fn,
)

optim = Adam(model.parameters(), lr = 1e-3)
for epoch in range(300):
    for idx, batch in enumerate(dataloader):
        imgs, labels = batch
        #imgs, labels = imgs.to(device), labels.to(device)
        imgs, labels = imgs.cuda(device=device_ids[0]), labels.cuda(device=device_ids[0])
        out, loss = model(imgs, labels)
        optim.zero_grad()
        loss.backward()
        optim.step()
        print("ep{},batch{}".format(epoch, idx), loss, )
    if epoch%5 == 0:
        print("save checkpoint")
        torch.save(model.state_dict(), "./ep{}_yolo_sub.ckpt".format(epoch))